## 使用mnist数据集做手写数字图片的识别训练
import numpy
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras import optimizers, layers, activations, losses, metrics

# 引入mnist数据
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
assert train_images.shape == (60000, 28, 28)
assert train_labels.shape == (60000,)
assert test_images.shape == (10000, 28, 28)
assert test_labels.shape == (10000,)

# 训练数据变形，3维变2维
train_images = train_images.reshape((60000, 28 * 28))
# 原始灰度图每个像素是uint8类型,取值范围 0-255
# 改为float32类型，取值范围改为 0-1
train_images = train_images.astype(dtype="float32") / 255

print(train_images.dtype)
print(train_images.shape)
print(train_images.size)
print(train_images.ndim)

# 定义一个由两个层组成的模型
model = tf.keras.Sequential(
    layers=[layers.Dense(512, activation=activations.relu),
            layers.Dense(10, activation=activations.softmax)])

# 编译模型
model.compile(optimizer=optimizers.RMSprop(),
              loss=losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])

# 训练模型
model.fit(train_images, train_labels, 128, 5)

# 测试数据需要和训练数据形状一致
test_images = test_images.reshape((10000, 28 * 28))
test_images = test_images.astype('float32') / 255

# shape(10000,10)
prediction = model.predict(test_images)
print(prediction.shape)
# shape(10000,1)
predict_labels = numpy.argmax(prediction, axis=1)
matches = test_labels == predict_labels
print(f'accuracy={matches.mean():.2f}')